import numpy as np
import os

dtype_emu = {
    np.float16: 1,
    np.float32: 2,
    np.int8: 3,
    np.int16: 4,
    np.int32: 5,
    np.complex64: 16,
    np.complex128: 17
}

def gen_isnan_data(
    target_shape=(17, 1023), 
    dtype=np.float32  # 支持：浮点/复数/整数类型
):
    # 验证数据类型：仅允许数值类型
    if not np.issubdtype(dtype, np.number):
        raise ValueError(f"仅支持数值类型（整数/浮点/复数），当前输入：{dtype}")
    
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)

    input_data = None
    if np.issubdtype(dtype, np.floating):
        # 浮点类型：生成随机值并插入5% NaN
        input_data = np.random.uniform(1, 100, target_shape).astype(dtype)
        nan_mask = np.random.choice([False, True], size=target_shape, p=[0.5, 0.5])
        input_data[nan_mask] = np.nan

    elif np.issubdtype(dtype, np.complexfloating):
        # 复数类型：实部和虚部使用对应精度的浮点类型（修复核心错误）
        # 例如：complex64 → 实部/虚部为float32；complex128 → 实部/虚部为float64
        float_dtype = np.dtype(dtype).kind.replace('c', 'f')  # 'c' → 'f'
        float_dtype = np.dtype(f"{float_dtype}{np.dtype(dtype).itemsize // 2}")
        
        real_part = np.random.uniform(1, 100, target_shape).astype(float_dtype)
        imag_part = np.random.uniform(1, 100, target_shape).astype(float_dtype)
        
        # 实部/虚部分别插入NaN
        real_nan_mask = np.random.choice([False, True], size=target_shape, p=[0.8, 0.2])
        real_part[real_nan_mask] = np.nan
        imag_nan_mask = np.random.choice([False, True], size=target_shape, p=[0.8, 0.2])
        imag_part[imag_nan_mask] = np.nan
        
        input_data = real_part + 1j * imag_part  # 组合为复数

    elif np.issubdtype(dtype, np.integer):
        # 整数类型：无NaN，生成随机整数
        input_data = np.random.randint(1, 100, target_shape).astype(dtype)

    # 计算golden结果
    golden = np.isnan(input_data)

    tiling = np.array([target_shape[0] * target_shape[1], dtype_emu[dtype]], dtype=np.uint32)

    # 保存数据
    tiling.tofile("./input/input_tiling.bin")
    input_data.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")

if __name__ == "__main__":
    # 示例1：浮点类型
    gen_isnan_data(target_shape=(17,1023), dtype=np.float32)
    
    # 示例2：复数类型
    # gen_isnan_data(target_shape=(17,1023), dtype=np.complex64)
    # gen_isnan_data(target_shape=(17,1023), dtype=np.complex128)
    
    # 示例3：整数类型
    # gen_isnan_data(target_shape=(17,1023), dtype=np.int32)
